package cn.juque.lucenecandy.helper;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.ArrayUtil;
import cn.juque.common.base.PageInfo;
import cn.juque.common.constants.MessageEnum;
import cn.juque.common.exception.AppException;
import cn.juque.lucenecandy.cache.IndexReaderCache;
import cn.juque.lucenecandy.cache.IndexWriterCache;
import lombok.extern.slf4j.Slf4j;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexNotFoundException;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.search.*;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Set;

/**
 * @author juque
 * @version 1.0.0
 * <ul>
 *     <li>LuceneHelper</li>
 * </ul>
 * @date 2023-12-17 01:18:21
 **/
@Slf4j
@Component("luceneHelper")
public class LuceneHelper {

    @Resource
    private IndexWriterCache indexWriterCache;

    @Resource
    private IndexReaderCache indexReaderCache;

    public List<Document> search(String index, BooleanQuery.Builder builder, Sort sort, Set<String> fieldToLoad) {
        Assert.notNull(index);
        Assert.notNull(builder);
        Assert.notNull(sort);
        try {
            long start = System.currentTimeMillis();
            IndexReader indexReader = this.indexReaderCache.get(index, false);
            IndexSearcher indexSearcher = new IndexSearcher(indexReader);
            int size = 10000;
            TopDocs topDocs;
            ScoreDoc lastDoc = null;
            List<Document> result = CollUtil.newArrayList();
            do {
                topDocs = indexSearcher.searchAfter(lastDoc, builder.build(), size, sort);
                if (topDocs.scoreDocs.length > 0) {
                    lastDoc = topDocs.scoreDocs[topDocs.scoreDocs.length - 1];
                    result.addAll(this.toDocument(indexSearcher, topDocs, fieldToLoad));
                }
            } while (topDocs.scoreDocs.length > 0);
            long cost = System.currentTimeMillis() - start;
            log.debug("search index:{} query:{}, size:{}, cost:{}ms", index, builder.build(), CollUtil.size(result), cost);
            return result;
        } catch (IndexNotFoundException e) {
            return new ArrayList<>(0);
        } catch (Exception e) {
            log.error("search index:{} 查询异常", index, e);
        }
        throw new AppException(MessageEnum.SYSTEM_ERROR);
    }

    public List<Document> searchByPage(String index, BooleanQuery.Builder builder, Sort sort, PageInfo pageInfo, Set<String> fieldToLoad) {
        long start = System.currentTimeMillis();
        List<Document> docList;
        try {
            IndexReader indexReader = this.indexReaderCache.get(index, false);
            IndexSearcher indexSearcher = new IndexSearcher(indexReader);
            // 查询出前一页最后一条记录，然后再查当前页
            if (pageInfo.getPage() <= 1) {
                TopDocs topDocs = indexSearcher.search(builder.build(), pageInfo.getLimit(), sort);
                pageInfo.setTotal(topDocs.totalHits.value);
                docList = this.toDocument(indexSearcher, topDocs, fieldToLoad);
            } else {
                int size = pageInfo.getLimit() * (pageInfo.getPage() - 1);
                TopDocs topDocs = indexSearcher.search(builder.build(), size, sort);
                ScoreDoc lastDoc = topDocs.scoreDocs[topDocs.scoreDocs.length - 1];
                topDocs = indexSearcher.searchAfter(lastDoc, builder.build(), pageInfo.getLimit(), sort);
                pageInfo.setTotal(topDocs.totalHits.value);
                docList = this.toDocument(indexSearcher, topDocs, fieldToLoad);
            }
            long end = System.currentTimeMillis() - start;
            if (end > 0) {
                log.debug("searchByPage index:{} query:{}, size:{}, cost time:{}ms", index, builder.build(), pageInfo.getTotal(), end);
            }
            return docList;
        } catch (IndexNotFoundException e) {
            return new ArrayList<>(0);
        } catch (Exception e) {
            log.error(" index:{} document searchByPage error", index, e);
        }
        throw new AppException(MessageEnum.SYSTEM_ERROR);
    }

    public void addDocuments(String index, List<Document> docList) {
        if (CollUtil.isEmpty(docList)) {
            return;
        }
        IndexWriter indexWriter = this.indexWriterCache.get(index, false);
        try {
            indexWriter.addDocuments(docList);
        } catch (Exception e) {
            log.error(" index:{} add documents error", index, e);
            throw new AppException(MessageEnum.SYSTEM_ERROR);
        }
    }

    public void updateByQuery(String index, Query query, Document document) {
        IndexWriter indexWriter = this.indexWriterCache.get(index, false);
        try {
            // 删除
            indexWriter.deleteDocuments(query);
            // 新增
            indexWriter.addDocument(document);
        } catch (Exception e) {
            log.error(" index:{} update document error", index, e);
            throw new AppException(MessageEnum.SYSTEM_ERROR);
        }
    }

    public void deleteDocuments(String index, Query... query) {
        IndexWriter indexWriter = this.indexWriterCache.get(index, false);
        try {
            indexWriter.deleteDocuments(query);
        } catch (Exception e) {
            log.error(" index:{} delete document error", index, e);
            throw new AppException(MessageEnum.SYSTEM_ERROR);
        }
    }

    public List<Document> toDocument(IndexSearcher searcher, TopDocs topDocs, Set<String> fieldToLoad) throws IOException {
        if (Objects.isNull(topDocs) || ArrayUtil.isEmpty(topDocs.scoreDocs)) {
            return new ArrayList<>(0);
        }
        ScoreDoc[] docs = topDocs.scoreDocs;
        List<Document> result = new ArrayList<>(docs.length);
        for (ScoreDoc doc : docs) {
            result.add(CollUtil.isEmpty(fieldToLoad) ? searcher.doc(doc.doc) : searcher.doc(doc.doc, fieldToLoad));
        }
        return result;
    }
}
